import numpy as np
import pywt
from scipy import signal
from scipy.signal import medfilt, butter, filtfilt

# TODO: 数据预处理
# 定义中值滤波函数
def median_filter(data):
    fs = 360  # 采样频率为250Hz
    tmp = data - medfilt(data, kernel_size=int(fs * 0.2) - 1)
    med_data = tmp - medfilt(tmp, kernel_size=int(fs * 0.6) - 1)
    return med_data


# 定义低通滤波函数
def low_pass(data):
    data = np.convolve(data, np.ones((5,)) / 5, mode='same')
    fs = 360  # 采样频率为250Hz
    order = 12
    cut_off = 35
    nyquist = 0.5 * fs
    normal_cutoff = cut_off / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    low_data = filtfilt(b, a, data)
    return low_data


# 定义小波变换函数定义
def DWT(ecg):
    w = pywt.Wavelet('db8')  # 选用Daubechies8小波
    maxlev = pywt.dwt_max_level(len(ecg), w.dec_len)
    threshold = 0.2  # Threshold for filtering
    coeffs = pywt.wavedec(ecg, 'db8', level=maxlev)  # 将信号进行小波分解
    # 这里就是对每一层的coffe进行更改
    for i in range(1, len(coeffs)):
        coeffs[i] = pywt.threshold(coeffs[i], threshold * max(coeffs[i]))

    datarec = pywt.waverec(coeffs, 'db8')  # 将信号进行小波重构
    return datarec


# 五点平滑滤波
def fliter(ecg):
    # 五点平滑滤波
    # ecg = np.convolve(ecg, np.ones((5,)) / 5, mode='same')
    # 工频干扰去除
    fs = 360  # 采样频率
    f0 = 50.0  # 要去除的工频干扰，50.0(抑制50Hz干扰)
    Q = 30.0  # 品质因数
    # 二阶IIR陷波数字滤波器，主要用于抑制信号中特定频率的干扰（如工频噪声）
    # w0，要抑制的中心频率（单位与 fs相同）；	品质因数，Q值越大，抑制带宽越窄； fs采样频率（必须提供，以明确频率单位）
    b, a = signal.iirnotch(w0=f0, Q=Q, fs=fs)
    # 工频干扰去除
    # 双向滤波函数，它通过正向和反向两次应用滤波器来实现零相位偏移，从而保持滤波后信号的相位特性不受影响
    ecg = signal.filtfilt(b, a, ecg)
    # 高通滤波器去除基线漂移
    [b, a] = signal.butter(3, [1 / 235 * 2, 20 / 235 * 2], 'bandpass')
    ecg = signal.filtfilt(b, a, ecg)
    return ecg
